import cv2
import os
import argparse
from datetime import datetime
import sys

def capture_faces(name, num_images=50, camera_id=0, output_dir="known_faces"):
    """
    捕获人脸图像用于训练人脸识别模型
    
    Args:
        name (str): 人员姓名，将用作保存图像文件名的前缀
        num_images (int): 要捕获的图像数量，默认为50张
        camera_id (int): 摄像头设备ID，默认为0（第一个摄像头）
        output_dir (str): 图像保存的目录路径，默认为"known_faces"
    """
    # 确保name参数使用正确的编码，特别是处理中文字符时
    if sys.platform.startswith('win'):
        # 在Windows系统上确保正确处理中文字符
        try:
            # 尝试用系统默认编码处理
            name = name.encode('gbk').decode('gbk')
        except UnicodeError:
            try:
                # 如果失败，尝试用UTF-8编码
                name = name.encode('utf-8').decode('utf-8')
            except UnicodeError:
                # 如果还失败，使用ASCII编码安全的字符串
                name = name.encode('ascii', 'ignore').decode('ascii')
    
    print(f"开始捕获人脸图像: {name}")
    print(f"目标数量: {num_images} 张")
    print(f"输出目录: {output_dir}")
    print("=" * 50)
    
    # 创建输出目录，如果目录不存在则创建
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        print(f"创建目录: {output_dir}")
    
    # 初始化摄像头设备
    cap = cv2.VideoCapture(camera_id)
    if not cap.isOpened():
        print(f"错误: 无法打开摄像头 {camera_id}")
        # 尝试其他设备ID (1-4)，在某些系统中摄像头可能分配了不同的ID
        for i in range(1, 5):
            print(f"尝试摄像头 {i}...")
            cap = cv2.VideoCapture(i)
            if cap.isOpened():
                print(f"成功打开摄像头 {i}")
                camera_id = i
                break
        if not cap.isOpened():
            print("无法打开任何摄像头")
            return
    
    # 设置摄像头分辨率，提高图像质量
    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
    
    # 加载OpenCV的人脸检测器（基于Haar特征的级联分类器）
    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
    
    # 初始化计数器
    # captured_count: 已成功捕获并保存的图像数量
    # frame_count: 处理的视频帧总数，用于控制自动捕获频率
    captured_count = 0
    frame_count = 0
    
    print("\n操作说明:")
    print("- 确保人脸正对摄像头")
    print("- 保持不同角度和表情")
    print("- 确保光线充足且均匀")
    print("- 按 'q' 或 'ESC' 键退出")
    print("- 按 'c' 键捕获当前帧中的人脸")
    print("- 自动捕获模式下会自动保存检测到的人脸")
    print("\n准备开始...")
    
    # 等待2秒让用户准备
    cv2.waitKey(2000)
    
    try:
        # 主循环：持续从摄像头读取帧并检测人脸
        while True:
            # 从摄像头读取一帧图像
            ret, frame = cap.read()
            if not ret:
                print("无法获取摄像头画面")
                break
            
            # 帧计数器递增
            frame_count += 1
            
            # 创建帧副本用于显示，避免在原始帧上绘制影响人脸检测
            display_frame = frame.copy()
            
            # 转换为灰度图像用于人脸检测（提高检测效率）
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # 检测人脸
            # scaleFactor: 图像缩放比例，值越小检测越细致但速度越慢
            # minNeighbors: 每个人脸候选矩形应保留的邻居数，值越大误检越少
            # minSize: 检测窗口的最小尺寸
            faces = face_cascade.detectMultiScale(
                gray,
                scaleFactor=1.1,
                minNeighbors=5,
                minSize=(30, 30)
            )
            
            # 在显示帧上绘制绿色矩形框标记检测到的人脸
            for (x, y, w, h) in faces:
                cv2.rectangle(display_frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
            
            # 在显示帧上添加文字信息
            cv2.putText(display_frame, f"已捕获: {captured_count}/{num_images}", 
                       (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            cv2.putText(display_frame, f"按 'c' 捕获, 'q' 退出", 
                       (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2)
            
            # 显示处理后的帧
            cv2.imshow('人脸捕获', display_frame)
            
            # 自动捕获模式（每30帧检查一次，且至少检测到一个人脸）
            # 这样可以避免保存过多重复的人脸图像
            if frame_count % 30 == 0 and len(faces) > 0 and captured_count < num_images:
                # 遍历所有检测到的人脸
                for i, (x, y, w, h) in enumerate(faces):
                    # 如果已达到目标数量则停止捕获
                    if captured_count >= num_images:
                        break
                        
                    # 提取人脸区域（裁剪出人脸部分）
                    face_img = frame[y:y+h, x:x+w]
                    
                    # 生成文件名，格式为姓名+编号，如GUO01.jpg, GUO02.jpg
                    filename = f"{name.upper()}{captured_count+1:02d}.jpg"
                    # 确保文件名在Windows系统上正确编码
                    if sys.platform.startswith('win'):
                        try:
                            filename = filename.encode('gbk').decode('gbk')
                        except UnicodeError:
                            filename = filename.encode('utf-8', 'ignore').decode('utf-8')
                    filepath = os.path.join(output_dir, filename)
                    
                    # 保存图像到指定目录
                    success = cv2.imwrite(filepath, face_img)
                    if success:
                        captured_count += 1
                        print(f"已保存: {filename} (共{captured_count}张)")
                    else:
                        print(f"保存失败: {filename}")
            
            # 等待按键输入（1毫秒）
            key = cv2.waitKey(1) & 0xFF
            
            # 按 'q' 或 ESC 键退出程序
            if key == ord('q') or key == 27:  # 27 is ESC key
                break
            
            # 按 'c' 键手动捕获当前帧中的人脸（立即保存当前检测到的人脸）
            elif key == ord('c') and len(faces) > 0:
                # 遍历所有检测到的人脸
                for i, (x, y, w, h) in enumerate(faces):
                    # 如果已达到目标数量则停止捕获
                    if captured_count >= num_images:
                        break
                        
                    # 提取人脸区域
                    face_img = frame[y:y+h, x:x+w]
                    
                    # 生成文件名
                    filename = f"{name.upper()}{captured_count+1:02d}.jpg"
                    # 确保文件名在Windows系统上正确编码
                    if sys.platform.startswith('win'):
                        try:
                            filename = filename.encode('gbk').decode('gbk')
                        except UnicodeError:
                            filename = filename.encode('utf-8', 'ignore').decode('utf-8')
                    filepath = os.path.join(output_dir, filename)
                    
                    # 保存图像
                    success = cv2.imwrite(filepath, face_img)
                    if success:
                        captured_count += 1
                        print(f"已保存: {filename} (共{captured_count}张)")
                    else:
                        print(f"保存失败: {filename}")
            
            # 检查是否达到目标数量
            if captured_count >= num_images:
                print(f"\n已完成! 成功捕获 {captured_count} 张人脸图像")
                break
    
    except KeyboardInterrupt:
        # 处理用户按Ctrl+C中断程序的情况
        print("\n用户中断程序")
    finally:
        # 释放摄像头资源并关闭所有OpenCV窗口
        cap.release()
        cv2.destroyAllWindows()
    
    # 输出最终统计信息
    print(f"\n捕获完成!")
    print(f"目标数量: {num_images}")
    print(f"实际捕获: {captured_count}")
    print(f"保存位置: {os.path.abspath(output_dir)}")

def list_captured_faces(output_dir="known_faces"):
    """
    列出指定目录中已捕获的人脸图像文件
    
    Args:
        output_dir (str): 要列出图像的目录路径，默认为"known_faces"
    """
    # 检查目录是否存在
    if not os.path.exists(output_dir):
        print(f"目录 {output_dir} 不存在")
        return
    
    # 获取目录中所有图像文件（jpg, png, jpeg格式）
    files = [f for f in os.listdir(output_dir) if f.endswith(('.jpg', '.png', '.jpeg'))]
    if not files:
        print(f"目录 {output_dir} 中没有图像文件")
        return
    
    # 打印所有图像文件名
    print(f"\n{output_dir} 中的图像文件:")
    for file in files:
        print(f"  - {file}")
    
    # 输出文件总数
    print(f"总计: {len(files)} 个文件")

def get_user_input():
    """
    获取用户输入的姓名和捕获数量
    
    Returns:
        tuple: (姓名, 数量) 或 (None, None)（输入格式错误时）
    """
    # 提示用户输入姓名和数量
    user_input = input("请输入姓名，并指定捕获数量。示例格式：GUO+50\n")
    try:
        # 解析用户输入，按"+"分割姓名和数量
        name, count_str = user_input.split('+')
        count = int(count_str)
        return name.strip(), count
    except ValueError:
        # 处理输入格式错误的情况
        print("输入格式错误，请使用格式：姓名+数量，例如：GUO+50")
        return None, None

def main():
    """
    主函数：解析命令行参数并调用相应功能
    """
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='人脸图像捕获工具')
    # 添加各种命令行参数
    parser.add_argument('--name', help='人员姓名（用作文件名前缀）')
    parser.add_argument('--count', type=int, default=50, help='要捕获的图像数量 (默认: 50)')
    parser.add_argument('--camera', type=int, default=0, help='摄像头ID (默认: 0)')
    parser.add_argument('--output', help='输出目录')
    parser.add_argument('--list', action='store_true', help='列出已捕获的图像')

    # 解析命令行参数
    args = parser.parse_args()
    
    # 如果用户指定了--list参数，则列出已捕获的图像
    if args.list:
        output_dir = args.output if args.output else "known_faces"
        list_captured_faces(output_dir)
    else:
        # 如果没有通过命令行参数提供姓名，则通过交互式输入获取
        if not args.name:
            name, count = get_user_input()
            if name is None or count is None:
                return
            args.name = name
            args.count = count
        
        # 设置输出目录为 data/name_face 格式，如 data/GUO_face/
        if not args.output:
            args.output = os.path.join("data", f"{args.name}_face")
        
        # 调用捕获人脸图像的函数
        capture_faces(args.name, args.count, args.camera, args.output)

if __name__ == "__main__":
    # 程序入口点
    main()